"""
Author: 'silencesoup' 'silencesoup@outlook.com'
Date: 2024-12-01 16:27:55
LastEditors: 'silencesoup' 'silencesoup@outlook.com'
LastEditTime: 2024-12-03 11:25:18
FilePath: \neimeng_2024\test.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
"""

from main import PastDataSource, model_infer
import xarray as xr
from loguru import logger


def test_read_attr_xrdataset():
    data_source = PastDataSource()
    gage_id_lst = ["NM000001", "NM000001"]
    data = data_source.read_attr_xrdataset(gage_id_lst, ["pre_mm_syr"])
    print(data)


def test_read_ts_xrdataset():
    data_source = PastDataSource()
    gage_id_lst = ["NM000001", "NM000002"]
    t_range = ["2024-10-01-01", "2024-10-11-01"]
    data = data_source.read_ts_xrdataset(
        gage_id_lst, t_range, var_lst=["gpm_tp", "sm_surface"], time_units=["3h"]
    )
    print(data)


def test_data_nc():
    data = xr.open_dataset("data/data.nc")
    print(data)


def test_model_infer():
    full_data = {
        "time_start": "2024-07-01-01",
        "gage_id_lst": ["NM000001"],
        "data": [
            {"basin": "NM000001", "time": "2024-07-01T01:00:00", "tp": 10.5},
            {"basin": "NM000001", "time": "2024-07-02T04:00:00", "tp": 15.0},
        ],
    }
    time_start = full_data["time_start"]
    gage_id_lst = full_data["gage_id_lst"]
    new_data = full_data["data"]
    pred = model_infer(
        time_start=time_start,
        gage_id_lst=gage_id_lst,
        new_data=new_data,
    )
    logger.warning(pred)


def test_read_mean_prcp():
    data_source = PastDataSource()
    gage_id_lst = ["NM000001"]
    data = data_source.read_mean_prcp(gage_id_lst, "mm/3h")
    logger.warning(data)


if __name__ == "__main__":
    # test_read_attr_xrdataset()
    # test_read_ts_xrdataset()
    # test_read_mean_prcp()
    # test_data_nc()
    test_model_infer()
